{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "在前面的RNN中，在训练时，如果我们输入的序列长度是5，那么输出的序列长度也是5，而Seq2Seq则输入和输出的长度可以是不一样的，因此非常适合做文本翻译的任务。",
   "id": "e25ee3b07acf5d94"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-23T14:46:30.388589Z",
     "start_time": "2025-05-23T14:46:22.043028Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "from datasets import load_dataset\n",
    "from torch.onnx.symbolic_opset9 import tensor\n",
    "\n",
    "# 加载全部数据分割\n",
    "full_dataset = load_dataset(\"Helsinki-NLP/opus-100\", \"en-zh\")\n",
    "\n",
    "train_data = full_dataset[\"train\"]\n",
    "valid_data = full_dataset[\"validation\"]\n",
    "test_data = full_dataset[\"test\"]\n",
    "\n",
    "print(f\"训练集样本数: {len(train_data)}\")\n",
    "print(f\"验证集样本数: {len(valid_data)}\")\n",
    "print(f\"测试集样本数: {len(test_data)}\")"
   ],
   "id": "9a7523062a08c226",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "训练集样本数: 1000000\n",
      "验证集样本数: 2000\n",
      "测试集样本数: 2000\n"
     ]
    }
   ],
   "execution_count": 149
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-23T14:53:53.971656Z",
     "start_time": "2025-05-23T14:53:53.966193Z"
    }
   },
   "cell_type": "code",
   "source": "train_data[:10]",
   "id": "4333c824dd487f82",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'translation': [{'en': 'Sixty-first session', 'zh': '第六十一届会议'},\n",
       "  {'en': 'I took some medicine for my mu for my mu my muscular for my muscular...',\n",
       "   'zh': '减轻酸... 酸痛的药 减轻酸痛的药'},\n",
       "  {'en': \"It's a challenge. God is challenging you. He's calling you a chump.\",\n",
       "   'zh': '上帝在挑战你，他说你是笨蛋'},\n",
       "  {'en': 'Oh, baby.', 'zh': '.. 寶貝'},\n",
       "  {'en': '- Lucinda?', 'zh': '- 盧辛達？'},\n",
       "  {'en': 'Introduction', 'zh': '一. 导言'},\n",
       "  {'en': 'Eric Topol: The wireless future of medicine',\n",
       "   'zh': 'Eric Topol：未来医疗的无线化'},\n",
       "  {'en': 'Damn it, Vaughn!', 'zh': '妈的, 沃恩!'},\n",
       "  {'en': 'General Recommendation XXII (Forty-ninth session, 1996)*',\n",
       "   'zh': '一般性建议二十一 (第四十八届会议， 1996年)*'},\n",
       "  {'en': '- Well, what about the witness at the bar?', 'zh': '在酒吧的那位证人呢'}]}"
      ]
     },
     "execution_count": 174,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 174
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-23T14:47:06.318189Z",
     "start_time": "2025-05-23T14:47:05.874939Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "# 加载预训练的分词器\n",
    "model_name = \"Helsinki-NLP/opus-mt-zh-en\"\n",
    "\n",
    "# Tokenizer代表分词器\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "# tokenizer添加一个特殊的标签\n",
    "tokenizer.add_special_tokens({\"additional_special_tokens\": [\"[SOS]\"]})\n",
    "\n",
    "print(tokenizer.convert_tokens_to_ids(\"[SOS]\"))\n",
    "print(tokenizer.eos_token)"
   ],
   "id": "204c9f69f1557d78",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "65001\n",
      "</s>\n"
     ]
    }
   ],
   "execution_count": 153
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-23T14:55:16.196524Z",
     "start_time": "2025-05-23T14:55:15.040260Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 定义预处理函数\n",
    "def preprocess_function(dataset):\n",
    "    # inputs = [data[\"zh\"] for data in dataset[\"translation\"]]\n",
    "    # targets = [data[\"en\"] for data in dataset[\"translation\"]]\n",
    "\n",
    "    inputs = [\"[SOS]\" + data[\"zh\"] for data in dataset[\"translation\"]]\n",
    "    targets = [\"[SOS]\" + data[\"en\"] for data in dataset[\"translation\"]]\n",
    "\n",
    "    # 输入最长128个字，不够的进行填充\n",
    "    model_inputs = tokenizer(inputs, max_length=128, truncation=True, padding=\"max_length\")\n",
    "\n",
    "    # 处理目标文本\n",
    "    labels = tokenizer(text_target=targets, max_length=128, truncation=True, padding=\"max_length\")\n",
    "\n",
    "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
    "    return model_inputs\n",
    "\n",
    "\n",
    "# 预处理\n",
    "tokenized_train_datasets = train_data.select(range(20)).map(preprocess_function, batched=True,\n",
    "                                                           remove_columns=[\"translation\"])\n",
    "tokenized_valid_datasets = valid_data.map(preprocess_function, batched=True, remove_columns=[\"translation\"])\n",
    "tokenized_test_datasets = test_data.map(preprocess_function, batched=True, remove_columns=[\"translation\"])\n",
    "\n",
    "# 查看处理后的数据格式\n",
    "# print(tokenized_train_datasets[:2])\n",
    "# print(train_data[:2])\n",
    "# print(len(tokenized_train_datasets))"
   ],
   "id": "48ec47fe0c16591b",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Map:   0%|          | 0/20 [00:00<?, ? examples/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "4ae6166d09ca4d8ea9d47e6ee8f50a28"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "execution_count": 186
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-23T14:55:17.906827Z",
     "start_time": "2025-05-23T14:55:17.903004Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 参数设置\n",
    "ZH_VOCAB_SIZE = tokenizer.vocab_size + 1\n",
    "EN_VOCAB_SIZE = tokenizer.vocab_size  +1\n",
    "HIDDEN_SIZE = 256\n",
    "BATCH_SIZE = 16\n",
    "LEARNING_RATE = 0.005"
   ],
   "id": "d2fff714a1743f7",
   "outputs": [],
   "execution_count": 187
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-23T14:55:19.190728Z",
     "start_time": "2025-05-23T14:55:19.182397Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader, TensorDataset\n",
    "\n",
    "dataset = TensorDataset(torch.tensor(tokenized_train_datasets['input_ids']),\n",
    "                        torch.tensor(tokenized_train_datasets['labels']))\n",
    "dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)"
   ],
   "id": "e14bb51bd9dfef07",
   "outputs": [],
   "execution_count": 188
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-23T14:56:34.892303Z",
     "start_time": "2025-05-23T14:55:20.950199Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from torch import nn, optim\n",
    "\n",
    "\n",
    "# 编码器\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(self, vocab_size, hidden_size):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(vocab_size, hidden_size)\n",
    "        self.rnn = nn.GRU(hidden_size, hidden_size, batch_first=True)\n",
    "\n",
    "    def forward(self, src):\n",
    "        embedded = self.embedding(src)\n",
    "        outputs, hidden = self.rnn(embedded)\n",
    "        return outputs, hidden\n",
    "\n",
    "\n",
    "# 解码器\n",
    "class Decoder(nn.Module):\n",
    "    def __init__(self, vocab_size, hidden_size):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(vocab_size, hidden_size)\n",
    "        self.rnn = nn.GRU(hidden_size, hidden_size, batch_first=True)\n",
    "        self.fc_out = nn.Linear(hidden_size, vocab_size)\n",
    "\n",
    "    def forward(self, input, hidden):\n",
    "        embedded = self.embedding(input)\n",
    "        outputs, hidden = self.rnn(embedded, hidden)\n",
    "        return self.fc_out(outputs), hidden\n",
    "\n",
    "\n",
    "# 训练配置\n",
    "encoder = Encoder(ZH_VOCAB_SIZE, HIDDEN_SIZE)\n",
    "decoder = Decoder(EN_VOCAB_SIZE, HIDDEN_SIZE)\n",
    "optimizer = optim.Adam(list(encoder.parameters()) + list(decoder.parameters()), lr=LEARNING_RATE)\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)\n",
    "\n",
    "\n",
    "# 训练函数\n",
    "def train(encoder, decoder, n_epochs=10):\n",
    "    for epoch in range(n_epochs):\n",
    "        for input, target in dataloader:\n",
    "            _, hidden = encoder(input)\n",
    "\n",
    "            # 准备解码器输入输出，其实这一步可以在数据处理时做掉\n",
    "            decoder_input = target[:, :-1]  # 移除最后一个token\n",
    "            decoder_target = target[:, 1:]  # 移除第一个token\n",
    "\n",
    "            decoder_output, _ = decoder(decoder_input, hidden)\n",
    "\n",
    "            # 计算损失\n",
    "            loss = criterion(\n",
    "                # decoder_output本来是(batch_size, seq_len, vocab_size)，变成(batch_size * seq_len, vocab_size)\n",
    "                # decoder_target本来是(batch_size, seq_len)，变成(batch_size * seq_len)\n",
    "                # reshape(-1)表示保留最后一个维度\n",
    "                decoder_output.reshape(-1, decoder_output.size(-1)),\n",
    "                decoder_target.reshape(-1)\n",
    "            )\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            print(f'Epoch {epoch + 1}, Loss: {loss:.4f}')\n",
    "\n",
    "\n",
    "# 开始训练\n",
    "train(encoder, decoder, n_epochs=100)"
   ],
   "id": "e0bb3d119e9e8fc9",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Loss: 11.0943\n",
      "Epoch 1, Loss: 11.0242\n",
      "Epoch 2, Loss: 9.6759\n",
      "Epoch 2, Loss: 8.8996\n",
      "Epoch 3, Loss: 7.3541\n",
      "Epoch 3, Loss: 5.8641\n",
      "Epoch 4, Loss: 4.4243\n",
      "Epoch 4, Loss: 4.5601\n",
      "Epoch 5, Loss: 3.4304\n",
      "Epoch 5, Loss: 2.7507\n",
      "Epoch 6, Loss: 2.5785\n",
      "Epoch 6, Loss: 2.2240\n",
      "Epoch 7, Loss: 1.8866\n",
      "Epoch 7, Loss: 1.9352\n",
      "Epoch 8, Loss: 1.4454\n",
      "Epoch 8, Loss: 1.7396\n",
      "Epoch 9, Loss: 1.1980\n",
      "Epoch 9, Loss: 1.0136\n",
      "Epoch 10, Loss: 0.9336\n",
      "Epoch 10, Loss: 0.8508\n",
      "Epoch 11, Loss: 0.7145\n",
      "Epoch 11, Loss: 0.9144\n",
      "Epoch 12, Loss: 0.6729\n",
      "Epoch 12, Loss: 0.5458\n",
      "Epoch 13, Loss: 0.5611\n",
      "Epoch 13, Loss: 0.6128\n",
      "Epoch 14, Loss: 0.4955\n",
      "Epoch 14, Loss: 0.6748\n",
      "Epoch 15, Loss: 0.4653\n",
      "Epoch 15, Loss: 0.6472\n",
      "Epoch 16, Loss: 0.4204\n",
      "Epoch 16, Loss: 0.5653\n",
      "Epoch 17, Loss: 0.4175\n",
      "Epoch 17, Loss: 0.4626\n",
      "Epoch 18, Loss: 0.4185\n",
      "Epoch 18, Loss: 0.3718\n",
      "Epoch 19, Loss: 0.3536\n",
      "Epoch 19, Loss: 0.5425\n",
      "Epoch 20, Loss: 0.3697\n",
      "Epoch 20, Loss: 0.3138\n",
      "Epoch 21, Loss: 0.3217\n",
      "Epoch 21, Loss: 0.3389\n",
      "Epoch 22, Loss: 0.3093\n",
      "Epoch 22, Loss: 0.3024\n",
      "Epoch 23, Loss: 0.2824\n",
      "Epoch 23, Loss: 0.3027\n",
      "Epoch 24, Loss: 0.2779\n",
      "Epoch 24, Loss: 0.2195\n",
      "Epoch 25, Loss: 0.2505\n",
      "Epoch 25, Loss: 0.2325\n",
      "Epoch 26, Loss: 0.2870\n",
      "Epoch 26, Loss: 0.1323\n",
      "Epoch 27, Loss: 0.2131\n",
      "Epoch 27, Loss: 0.4752\n",
      "Epoch 28, Loss: 0.2133\n",
      "Epoch 28, Loss: 0.2335\n",
      "Epoch 29, Loss: 0.1810\n",
      "Epoch 29, Loss: 0.3072\n",
      "Epoch 30, Loss: 0.2129\n",
      "Epoch 30, Loss: 0.1282\n",
      "Epoch 31, Loss: 0.1694\n",
      "Epoch 31, Loss: 0.1703\n",
      "Epoch 32, Loss: 0.1625\n",
      "Epoch 32, Loss: 0.1365\n",
      "Epoch 33, Loss: 0.1417\n",
      "Epoch 33, Loss: 0.1449\n",
      "Epoch 34, Loss: 0.1371\n",
      "Epoch 34, Loss: 0.1084\n",
      "Epoch 35, Loss: 0.1127\n",
      "Epoch 35, Loss: 0.1831\n",
      "Epoch 36, Loss: 0.1082\n",
      "Epoch 36, Loss: 0.1066\n",
      "Epoch 37, Loss: 0.1026\n",
      "Epoch 37, Loss: 0.0861\n",
      "Epoch 38, Loss: 0.0977\n",
      "Epoch 38, Loss: 0.0755\n",
      "Epoch 39, Loss: 0.0818\n",
      "Epoch 39, Loss: 0.0982\n",
      "Epoch 40, Loss: 0.0786\n",
      "Epoch 40, Loss: 0.0580\n",
      "Epoch 41, Loss: 0.0760\n",
      "Epoch 41, Loss: 0.0446\n",
      "Epoch 42, Loss: 0.0543\n",
      "Epoch 42, Loss: 0.0837\n",
      "Epoch 43, Loss: 0.0594\n",
      "Epoch 43, Loss: 0.0367\n",
      "Epoch 44, Loss: 0.0465\n",
      "Epoch 44, Loss: 0.0454\n",
      "Epoch 45, Loss: 0.0405\n",
      "Epoch 45, Loss: 0.0586\n",
      "Epoch 46, Loss: 0.0402\n",
      "Epoch 46, Loss: 0.0402\n",
      "Epoch 47, Loss: 0.0362\n",
      "Epoch 47, Loss: 0.0396\n",
      "Epoch 48, Loss: 0.0354\n",
      "Epoch 48, Loss: 0.0289\n",
      "Epoch 49, Loss: 0.0313\n",
      "Epoch 49, Loss: 0.0303\n",
      "Epoch 50, Loss: 0.0254\n",
      "Epoch 50, Loss: 0.0418\n",
      "Epoch 51, Loss: 0.0257\n",
      "Epoch 51, Loss: 0.0270\n",
      "Epoch 52, Loss: 0.0252\n",
      "Epoch 52, Loss: 0.0207\n",
      "Epoch 53, Loss: 0.0224\n",
      "Epoch 53, Loss: 0.0262\n",
      "Epoch 54, Loss: 0.0231\n",
      "Epoch 54, Loss: 0.0170\n",
      "Epoch 55, Loss: 0.0246\n",
      "Epoch 55, Loss: 0.0130\n",
      "Epoch 56, Loss: 0.0204\n",
      "Epoch 56, Loss: 0.0257\n",
      "Epoch 57, Loss: 0.0182\n",
      "Epoch 57, Loss: 0.0278\n",
      "Epoch 58, Loss: 0.0189\n",
      "Epoch 58, Loss: 0.0177\n",
      "Epoch 59, Loss: 0.0169\n",
      "Epoch 59, Loss: 0.0192\n",
      "Epoch 60, Loss: 0.0159\n",
      "Epoch 60, Loss: 0.0195\n",
      "Epoch 61, Loss: 0.0164\n",
      "Epoch 61, Loss: 0.0132\n",
      "Epoch 62, Loss: 0.0151\n",
      "Epoch 62, Loss: 0.0165\n",
      "Epoch 63, Loss: 0.0146\n",
      "Epoch 63, Loss: 0.0142\n",
      "Epoch 64, Loss: 0.0149\n",
      "Epoch 64, Loss: 0.0101\n",
      "Epoch 65, Loss: 0.0131\n",
      "Epoch 65, Loss: 0.0127\n",
      "Epoch 66, Loss: 0.0121\n",
      "Epoch 66, Loss: 0.0145\n",
      "Epoch 67, Loss: 0.0140\n",
      "Epoch 67, Loss: 0.0077\n",
      "Epoch 68, Loss: 0.0112\n",
      "Epoch 68, Loss: 0.0142\n",
      "Epoch 69, Loss: 0.0108\n",
      "Epoch 69, Loss: 0.0135\n",
      "Epoch 70, Loss: 0.0097\n",
      "Epoch 70, Loss: 0.0173\n",
      "Epoch 71, Loss: 0.0105\n",
      "Epoch 71, Loss: 0.0113\n",
      "Epoch 72, Loss: 0.0113\n",
      "Epoch 72, Loss: 0.0074\n",
      "Epoch 73, Loss: 0.0099\n",
      "Epoch 73, Loss: 0.0119\n",
      "Epoch 74, Loss: 0.0099\n",
      "Epoch 74, Loss: 0.0106\n",
      "Epoch 75, Loss: 0.0109\n",
      "Epoch 75, Loss: 0.0065\n",
      "Epoch 76, Loss: 0.0090\n",
      "Epoch 76, Loss: 0.0126\n",
      "Epoch 77, Loss: 0.0101\n",
      "Epoch 77, Loss: 0.0060\n",
      "Epoch 78, Loss: 0.0089\n",
      "Epoch 78, Loss: 0.0079\n",
      "Epoch 79, Loss: 0.0081\n",
      "Epoch 79, Loss: 0.0125\n",
      "Epoch 80, Loss: 0.0090\n",
      "Epoch 80, Loss: 0.0063\n",
      "Epoch 81, Loss: 0.0089\n",
      "Epoch 81, Loss: 0.0059\n",
      "Epoch 82, Loss: 0.0078\n",
      "Epoch 82, Loss: 0.0085\n",
      "Epoch 83, Loss: 0.0074\n",
      "Epoch 83, Loss: 0.0117\n",
      "Epoch 84, Loss: 0.0070\n",
      "Epoch 84, Loss: 0.0088\n",
      "Epoch 85, Loss: 0.0071\n",
      "Epoch 85, Loss: 0.0098\n",
      "Epoch 86, Loss: 0.0072\n",
      "Epoch 86, Loss: 0.0075\n",
      "Epoch 87, Loss: 0.0074\n",
      "Epoch 87, Loss: 0.0058\n",
      "Epoch 88, Loss: 0.0065\n",
      "Epoch 88, Loss: 0.0084\n",
      "Epoch 89, Loss: 0.0069\n",
      "Epoch 89, Loss: 0.0062\n",
      "Epoch 90, Loss: 0.0069\n",
      "Epoch 90, Loss: 0.0061\n",
      "Epoch 91, Loss: 0.0063\n",
      "Epoch 91, Loss: 0.0080\n",
      "Epoch 92, Loss: 0.0068\n",
      "Epoch 92, Loss: 0.0047\n",
      "Epoch 93, Loss: 0.0063\n",
      "Epoch 93, Loss: 0.0062\n",
      "Epoch 94, Loss: 0.0061\n",
      "Epoch 94, Loss: 0.0062\n",
      "Epoch 95, Loss: 0.0060\n",
      "Epoch 95, Loss: 0.0066\n",
      "Epoch 96, Loss: 0.0061\n",
      "Epoch 96, Loss: 0.0056\n",
      "Epoch 97, Loss: 0.0063\n",
      "Epoch 97, Loss: 0.0049\n",
      "Epoch 98, Loss: 0.0060\n",
      "Epoch 98, Loss: 0.0051\n",
      "Epoch 99, Loss: 0.0056\n",
      "Epoch 99, Loss: 0.0060\n",
      "Epoch 100, Loss: 0.0061\n",
      "Epoch 100, Loss: 0.0040\n"
     ]
    }
   ],
   "execution_count": 189
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-23T14:58:58.645508Z",
     "start_time": "2025-05-23T14:58:58.620646Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 翻译函数\n",
    "def translate(sentence, encoder, decoder):\n",
    "    model_inputs = tokenizer(sentence, max_length=128, truncation=True, padding=\"max_length\")\n",
    "\n",
    "    _, hidden = encoder(torch.tensor(model_inputs['input_ids']))\n",
    "    trg_indexes = [tokenizer.convert_tokens_to_ids(\"[SOS]\")]\n",
    "\n",
    "    for _ in range(50):\n",
    "        trg_tensor = torch.LongTensor([trg_indexes[-1]])\n",
    "        with torch.no_grad():\n",
    "            output, hidden = decoder(trg_tensor, hidden)\n",
    "        pred_token = output.argmax().item()\n",
    "        trg_indexes.append(pred_token)\n",
    "        if pred_token == tokenizer.eos_token_id:\n",
    "            break\n",
    "\n",
    "    return ' '.join([tokenizer.convert_ids_to_tokens(idx) for idx in trg_indexes[1:-1]])\n",
    "\n",
    "\n",
    "# 测试翻译\n",
    "test_sentence = \"第六十一届会议\"\n",
    "print(translate(test_sentence, encoder, decoder))"
   ],
   "id": "392d668e7d115e97",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "▁Sixty - first ▁session\n"
     ]
    }
   ],
   "execution_count": 193
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
